from storage.chroma_manager import ChromaManager
import gradio as gr
gr.analytic_enabled = False
manager = ChromaManager()


def add_document(collection_name, doc_id, content, metadata):
    try:
        manager = ChromaManager(collection_name=collection_name)
        if not doc_id.strip():
            return "错误：文档ID不能为空"
        if not content.strip():
            return "错误：文档内容不能为空"

        manager.add_document(doc_id, content, eval(metadata) if metadata else None)
        return f"文档 {doc_id} 添加成功"
    except Exception as e:
        return f"添加失败: {str(e)}"

# 新增修改文档函数
def update_document(collection_name, doc_id, new_content, new_metadata):
    try:
        manager = ChromaManager(collection_name=collection_name)
        manager.add_document(doc_id, new_content, eval(new_metadata) if new_metadata else None)
        return f"文档 {doc_id} 更新成功"
    except Exception as e:
        return f"更新失败: {str(e)}"

# 新增清空输入函数
def clear_add():
    return "", "", ""

def delete_document(collection_name, doc_id):
    try:
        manager = ChromaManager(collection_name=collection_name)
        manager.delete_document(doc_id)
        return f"文档 {doc_id} 删除成功"
    except Exception as e:
        return f"删除失败: {str(e)}"

def search_documents(collection_name, query):
    try:
        manager = ChromaManager(collection_name=collection_name)
        results = manager.search(query)
        merged_content = "\n\n\n\n".join(
            f"证据 {i + 1}（ID: {doc.id}）：\n{doc.page_content}"
            for i, (doc, _) in enumerate(results)
        )
        return merged_content
    except Exception as e:
        return f"{str(e)}"

# 新增文档查询函数
def query_document(collection_name, doc_id):
    try:
        manager = ChromaManager(collection_name=collection_name)
        result = manager.get_all_documents()
        if doc_id not in result['ids']:
            return None, None, "文档不存在"

        index = result['ids'].index(doc_id)
        content = result['documents'][index]
        metadata = str(result['metadatas'][index]) if result['metadatas'] else ""
        return content, metadata, "查询成功"
    except Exception as e:
        return None, None, f"查询失败: {str(e)}"


with gr.Blocks() as demo:
    gr.Markdown("## 知识库管理系统")

    # 新增全局设置区域
    with gr.Row():
        collection_name = gr.Textbox(
            label="集合名称",
            value="knowledge_base",
            placeholder="输入要操作的集合名称",
            interactive=True
        )

    with gr.Tab("添加文档"):
        with gr.Row():
            doc_id = gr.Textbox(label="文档ID*", placeholder="必填字段")  # 添加星号标记
            content = gr.Textbox(label="文档内容*", lines=3, placeholder="必填字段")
            metadata = gr.Textbox(label="元数据（字典格式）", placeholder="例如: {'source': 'manual'}")

        with gr.Row():
            add_btn = gr.Button("提交")
            clear_btn = gr.Button("清空")
        add_output = gr.Textbox(label="操作结果")
        clear_btn.click(clear_add, None, [doc_id, content, metadata])
        # 新增修改文档标签页

    with gr.Tab("修改文档"):
        with gr.Row():
            update_id = gr.Textbox(label="要修改的文档ID")
            query_btn = gr.Button("查询文档")
        with gr.Row():
            new_content = gr.Textbox(label="新内容", lines=3, interactive=True)
            new_metadata = gr.Textbox(label="新元数据", placeholder="例如: {'source': 'updated'}", interactive=True)
        with gr.Row():
            update_btn = gr.Button("提交修改")
            query_output = gr.Textbox(label="查询结果")

    # 新增交互绑定
    query_btn.click(
        query_document,
        inputs=[collection_name,update_id],
        outputs=[new_content, new_metadata, query_output]
    )
    update_btn.click(
        update_document,
        inputs=[collection_name,update_id, new_content, new_metadata],
        outputs=query_output
    )


    with gr.Tab("删除文档"):
        del_id = gr.Textbox(label="要删除的文档ID")
        del_btn = gr.Button("删除")
        del_output = gr.Textbox(label="操作结果")

    with gr.Tab("搜索文档"):
        search_query = gr.Textbox(label="搜索关键词")
        search_btn = gr.Button("搜索")
        search_output = gr.Textbox(label="结果", lines=6)

    add_btn.click(add_document, [collection_name, doc_id, content, metadata], add_output)
    del_btn.click(delete_document,[collection_name, del_id], del_output)
    search_btn.click(search_documents, [collection_name, search_query], search_output)

if __name__ == "__main__":
    demo.launch()
